package cz.semjob.learning.classifier;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.apache.log4j.Logger;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import weka.classifiers.Classifier;
import cz.semjob.document.annotation.LabeledRelationWordsPair;
import cz.semjob.document.annotation.LabeledWords;
import cz.semjob.document.annotation.RelationWordsPair;
import cz.semjob.document.annotation.metadata.AnnotationTypeFactory;
import cz.semjob.document.annotation.metadata.Entity;
import cz.semjob.document.annotation.metadata.Relation;
import cz.semjob.document.processing.DocumentInstance;
import cz.semjob.learning.instances.InstanceMapping;
import cz.semjob.learning.instances.UnlabeledRelationInstance;
import cz.semjob.learning.instances.creator.ITextInstanceCreator;
import cz.semjob.learning.instances.loader.InstancesCreator;
import cz.semjob.learning.instances.loader.InstancesLoadingException;
import cz.semjob.learning.relation.IRelationInstanceCreator;

@Component
public class ClassifierExecutor {

	@Autowired
	private InstancesCreator instancesLoader;

	@Autowired
	private AnnotationTypeFactory annotationTypeFactory;

	private static Logger logger = Logger.getLogger(ClassifierExecutor.class);

	public Map<DocumentInstance, List<LabeledWords>> executeClassificationForEntity(
			Classifier classifier, Entity entity,
			ITextInstanceCreator instanceCreator,
			List<DocumentInstance> documentInstances, Object... dataSources)
			throws ClassificationException {

		Map<DocumentInstance, List<InstanceMapping>> instances;
		try {
			instances = instancesLoader.createInstances(documentInstances,
					instanceCreator, dataSources);
		} catch (InstancesLoadingException e) {
			logger.error(e);
			throw new ClassificationException(e);
		}
		Map<DocumentInstance, List<LabeledWords>> foundAnnotations = new HashMap<DocumentInstance, List<LabeledWords>>();
		for (DocumentInstance documentInstance : instances.keySet()) {

			List<InstanceMapping> instanceMappings = instances
					.get(documentInstance);
			logger.debug("Candidate words count for document instance "
					+ documentInstance.getFileName() + " : "
					+ instanceMappings.size());
			List<LabeledWords> foundDocumentAnnotations = executeEntityClassificationOneDocumentInstance(
					documentInstance, instanceMappings, classifier, entity);
			logger.debug("Found words count for document instance "
					+ documentInstance.getFileName() + " and entity " + entity
					+ " : " + foundDocumentAnnotations.size());
			foundAnnotations.put(documentInstance, foundDocumentAnnotations);

		}
		return foundAnnotations;

	}

	public Map<DocumentInstance, Map<Relation, List<LabeledRelationWordsPair>>> executeClassificationForRelations(
			Classifier classifier, String relationType,
			IRelationInstanceCreator instanceCreator,
			Map<DocumentInstance, List<RelationWordsPair>> documentInstances)
			throws ClassificationException {

		Map<DocumentInstance, List<InstanceMapping>> instances;
		List<Relation> relations = annotationTypeFactory
				.getRelationsByType(relationType);
		try {
			instances = instancesLoader.createInstances(documentInstances,
					instanceCreator, relations);
		} catch (InstancesLoadingException e) {
			logger.error(e);
			throw new ClassificationException(e);
		}
		Map<DocumentInstance, Map<Relation, List<LabeledRelationWordsPair>>> foundAnnotations = new HashMap<DocumentInstance, Map<Relation, List<LabeledRelationWordsPair>>>();
		for (DocumentInstance documentInstance : instances.keySet()) {

			List<InstanceMapping> instanceMappings = instances
					.get(documentInstance);

			Map<Relation, List<LabeledRelationWordsPair>> foundDocumentAnnotations = executeRelationClassificationOneDocumentInstance(
					documentInstance, instanceMappings, classifier, relations);
			foundAnnotations.put(documentInstance, foundDocumentAnnotations);

		}
		return foundAnnotations;

	}

	public Map<DocumentInstance, List<LabeledWords>> executeClassification(
			Classifier classifier, Entity entity,
			ITextInstanceCreator instanceCreator, String inputPath,
			String processedPath, Object... dataSources)
			throws ClassificationException {

		logger.info("Classifying text documents in: " + inputPath);

		Map<DocumentInstance, List<InstanceMapping>> instances;
		try {
			instances = instancesLoader.loadInstances(inputPath, processedPath,
					instanceCreator, dataSources);
		} catch (InstancesLoadingException e) {
			logger.error(e);
			throw new ClassificationException(e);
		}
		Map<DocumentInstance, List<LabeledWords>> foundAnnotations = new HashMap<DocumentInstance, List<LabeledWords>>();

		for (DocumentInstance documentInstance : instances.keySet()) {

			List<InstanceMapping> instanceMappings = instances
					.get(documentInstance);

			List<LabeledWords> foundDocumentAnnotations = executeEntityClassificationOneDocumentInstance(
					documentInstance, instanceMappings, classifier, entity);
			foundAnnotations.put(documentInstance, foundDocumentAnnotations);

		}
		return foundAnnotations;
	}

	private List<LabeledWords> executeEntityClassificationOneDocumentInstance(
			DocumentInstance documentInstance,
			List<InstanceMapping> instanceMappings, Classifier classifier,
			Entity entity) throws ClassificationException {
		List<LabeledWords> foundDocumentAnnotations = new ArrayList<LabeledWords>();

		logger.debug("====== " + documentInstance.getFileName() + " ======");
		for (InstanceMapping instanceMapping : instanceMappings) {

			double isEntity;
			try {
				isEntity = classifier.classifyInstance(instanceMapping
						.getInstance());

				if (isEntity == 0d) {
					double[] distribution = classifier
							.distributionForInstance(instanceMapping
									.getInstance());
					foundDocumentAnnotations.add(new LabeledWords(entity,
							instanceMapping.getUnlabeledTextInstance()
									.getWords()));
					logger.debug(instanceMapping.getUnlabeledTextInstance()
							.getText() + ": " + Arrays.toString(distribution));
				}
			} catch (Exception e) {
				logger.error(e);
				throw new ClassificationException(e);
			}
		}
		return foundDocumentAnnotations;
	}

	private Map<Relation, List<LabeledRelationWordsPair>> executeRelationClassificationOneDocumentInstance(
			DocumentInstance documentInstance,
			List<InstanceMapping> instanceMappings, Classifier classifier,
			List<Relation> relations) throws ClassificationException {
		Map<Relation, List<LabeledRelationWordsPair>> foundDocumentAnnotations = new HashMap<Relation, List<LabeledRelationWordsPair>>();

		logger.debug("====== " + documentInstance.getFileName() + " ======");
		for (InstanceMapping instanceMapping : instanceMappings) {

			double relationIndex;
			try {
				relationIndex = classifier.classifyInstance(instanceMapping
						.getInstance());
			} catch (Exception e) {
				logger.error(e);
				throw new ClassificationException(e);
			}
			if (relationIndex < relations.size()) {
				Relation relation = relations.get((int) relationIndex);
				if (!relation.equals(Relation.getNoneRelation(relation
						.getType()))) {
					List<LabeledRelationWordsPair> relationWordsPairs = null;
					if (foundDocumentAnnotations.containsKey(relation)) {
						relationWordsPairs = foundDocumentAnnotations
								.get(relation);
					} else {
						relationWordsPairs = new ArrayList<LabeledRelationWordsPair>();
						foundDocumentAnnotations.put(relation,
								relationWordsPairs);
					}
					LabeledRelationWordsPair foundRelationWordsPair = new LabeledRelationWordsPair(
							relation,
							((UnlabeledRelationInstance) instanceMapping
									.getUnlabeledTextInstance())
									.getRelationWordsPair());
					// quick fix, adding only when relation is not already added
					// if (!relationWordsPairs.contains(foundRelationWordsPair))
					// {
					relationWordsPairs.add(foundRelationWordsPair);
					logger.debug(relation.getName() + ":"
							+ foundRelationWordsPair.toPrettyString());
					// } else {
					// logger.warn("Relation already added to classification results: "
					// + foundRelationWordsPair);
					// }
				}
			}
		}
		return foundDocumentAnnotations;
	}

	public void setInstancesLoader(InstancesCreator instancesLoader) {
		this.instancesLoader = instancesLoader;
	}

	public void setAnnotationTypeFactory(
			AnnotationTypeFactory annotationTypeFactory) {
		this.annotationTypeFactory = annotationTypeFactory;
	}

}
